package encrypt

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"io"
	"sync"
)

var (
	onceAes   sync.Once
	singleAes *Aes
)

type Aes struct {
	key []byte
}

// GetAes 获取AES实例
func GetAes(key string) *Aes {
	onceAes.Do(func() {
		singleAes = &Aes{[]byte(key)}
	})
	return singleAes
}

// Encrypt 加密
func (a Aes) Encrypt(plaintext string) string {
	plainBytes := []byte(plaintext)

	block, err := aes.NewCipher(a.key)
	if err != nil {
		panic(err)
	}
	// 密码块长度
	blockSize := block.BlockSize()
	plainBytes = pkcs7Padding(plainBytes, blockSize)

	ciphertext := make([]byte, aes.BlockSize+len(plainBytes))
	// 生成iv
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext[aes.BlockSize:], plainBytes)

	return base64.StdEncoding.EncodeToString(ciphertext)
}

// Decrypt 解密
func (a Aes) Decrypt(ciphertext string) string {
	cipherBytes, _ := base64.StdEncoding.DecodeString(ciphertext)
	block, err := aes.NewCipher(a.key)
	if err != nil {
		panic(err)
	}
	// 密码块长度
	iv := cipherBytes[:aes.BlockSize]
	cipherBytes = cipherBytes[aes.BlockSize:]
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(cipherBytes, cipherBytes)
	return string(pkcs7UnPadding(cipherBytes))
}

func pkcs7Padding(plainBytes []byte, blockSize int) []byte {
	padding := blockSize - len(plainBytes)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(plainBytes, padText...)
}

func pkcs7UnPadding(ciphertext []byte) []byte {
	length := len(ciphertext)
	unPadding := int(ciphertext[length-1])
	return ciphertext[:(length - unPadding)]
}
